import pandas as pd
import json
import os
import argparse

def normalize_path(p: str) -> str:
    """规范化路径方便匹配"""
    if not p:
        return ""
    p = str(p).strip()
    p = os.path.normpath(p).replace("\\", "/")
    if p.startswith("./"):
        p = p[2:]
    return p

def main():
    parser = argparse.ArgumentParser(description="将 JSON 的 selected_caption 写入 CSV 的 Caption 列")
    parser.add_argument("--csv_file", help="输入 CSV 文件路径")
    parser.add_argument("--json_file", help="输入 JSON 文件路径")
    parser.add_argument("--out_file", help="输出 CSV 文件路径")
    args = parser.parse_args()

    # 读 csv
    df = pd.read_csv(args.csv_file, dtype=str, keep_default_na=False)
    if "Source" not in df.columns:
        raise ValueError("CSV 文件缺少 Source 列")

    # 读 json
    with open(args.json_file, "r", encoding="utf-8") as f:
        data = json.load(f)

    # 建立映射表
    mapping = {normalize_path(d["image"]): d.get("selected_caption", "")
               for d in data if "image" in d}

    # 匹配写入 Caption
    df["Caption"] = df["Source"].apply(lambda x: mapping.get(normalize_path(x), ""))

    # 保存结果
    df.to_csv(args.out_file, index=False, encoding="utf-8-sig")
    print(f"处理完成，输出到 {args.out_file}")

if __name__ == "__main__":
    main()
